#include "selfupdate.h"

#include <wininet.h>
#include <tchar.h>
#include <wchar.h>
#include <stdio.h>

#pragma comment (lib, "wininet.lib")

#define SUI_HEADER "SUI!"
#define SUI_VERSION 0x0200

#define MEMORY2VALUE(l,t,r) (l) = *(t*)(r);

char* GetContentFromUrl(const wchar_t* Url, BOOL ToString, DWORD* Length)
{
#define MAXBLOCKSIZE (512*1024)

    char* Content = NULL;
    int ContentLength = 0;
    int Allocated = 0;

    HINTERNET hSession  =  InternetOpen( TEXT("IE"), INTERNET_OPEN_TYPE_PRECONFIG, NULL, NULL,  0 );

    if  (hSession  !=  NULL)
    {
        HINTERNET handle2  =  InternetOpenUrlW(hSession, Url, NULL,  0 , INTERNET_FLAG_DONT_CACHE,  0 );

        if  (handle2  !=  NULL)
        {
            while (1)
            {
                DWORD Read = 0;

                if (ContentLength + MAXBLOCKSIZE >= Allocated)
                {
                    Allocated += MAXBLOCKSIZE * 2;
                    Content = (char*)realloc(Content, Allocated);
                }

                InternetReadFile(handle2, Content + ContentLength, MAXBLOCKSIZE,  &Read);

                if (Read > 0)
                {
                    ContentLength += Read;
                }
                else
                {
                    if (ToString)
                    {
                        Content[ContentLength] = 0;
                        ContentLength++;
                    }

                    Content = (char*)realloc(Content, ContentLength);
                    break;
                }

            }

            InternetCloseHandle(handle2);
            handle2  =  NULL;
        } 

        InternetCloseHandle(hSession);
        hSession  =  NULL;
    } 

    if (!ContentLength)
    {
        HANDLE file = CreateFileW(Url, GENERIC_READ, FILE_SHARE_READ | FILE_SHARE_WRITE, NULL,
            OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);


        if (file == INVALID_HANDLE_VALUE)
        {
            file = NULL;
        }	

        if (file)
        {
            DWORD nRW = 0;
            ContentLength = GetFileSize(file, NULL);

            Content = (char*)malloc(ContentLength);

            if ((!ReadFile(file, Content, ContentLength, &nRW, NULL)) || (nRW == 0))
            {
                free(Content);
                Content = NULL;
                ContentLength = 0;
            }

            CloseHandle(file);
        }

    }

    if (Length)
        *Length = ContentLength;

    return Content;
#undef MAXBLOCKSIZE
}

void CheckUpdateInfo(SELFUPDATEPARAM* Param, const char* Str, DWORD Len)
{
#define ADVANCE_MEM2VAL(v, t) \
    if (Len < sizeof(t)) { v = -1; return;} \
    MEMORY2VALUE(v, t, Mem);\
    Len += sizeof(t);\
    Mem += sizeof(t);

#define ADVANCE_MEM2TEXT(v) \
    ADVANCE_MEM2VAL(ShortVar, WORD);\
    if (ShortVar < 0) return; \
    if (Len < ShortVar) return; \
    memcpy(v, Mem, ShortVar*2); \
    v[ShortVar] = 0; \
    Len -= ShortVar*2;\
    Mem += ShortVar*2;


    WORD ShortVar;
    LONG LongVar;

    char* Mem = (char*)Str;

    Param->UpdateFound = FALSE;

    ADVANCE_MEM2VAL(LongVar, LONG);

    if (0 != memcmp(&LongVar, SUI_HEADER, 4))
        return;

    ADVANCE_MEM2VAL(ShortVar, WORD);

    if (ShortVar < SUI_VERSION)
        return;

    if (ShortVar > SUI_VERSION)
    {
        Param->UpdateFound = TRUE;
        return;
    }

    ADVANCE_MEM2VAL(LongVar, DWORD);

    if (LongVar <= Param->VersionCode)
        return; // No need to update

    ADVANCE_MEM2TEXT(Param->Homepage);
    ADVANCE_MEM2TEXT(Param->UpdateFilePath);
    ADVANCE_MEM2TEXT(Param->UpdateFileName);
    ADVANCE_MEM2TEXT(Param->UpdateLogFileName);
    ADVANCE_MEM2TEXT(Param->DownloadPage);

    Param->UpdateFound = TRUE;

#undef ADVANCE_MEM2VAL
#undef ADVANCE_MEM2TEXT
}

void ProduceUpdateInfo(SELFUPDATEPARAM* Param)
{
    TCHAR Drive[_MAX_DRIVE];
    TCHAR Dir[_MAX_DIR];
    TCHAR Filename[MAX_PATH];
    FILE* File;

    WORD ShortVar = SUI_VERSION;

    DWORD Ret = GetModuleFileName(NULL, Filename, MAX_PATH);

    _tsplitpath(Filename, Drive, Dir, NULL,NULL);

    _stprintf(Filename,TEXT("%s%supdateinfox.sui"), Drive, Dir);

    File = _tfopen(Filename, TEXT("wb"));

    fwrite(SUI_HEADER, 4, 1, File);
    fwrite(&ShortVar, 2, 1, File);
    fwrite(&Param->VersionCode, sizeof(Param->VersionCode), 1, File);

#define WRITE_STRING(s) \
    ShortVar = wcslen(s);\
    fwrite(&ShortVar, 2, 1, File); \
    fwrite(s, ShortVar, 2, File);

    WRITE_STRING(Param->Homepage);
    WRITE_STRING(Param->UpdateFilePath);
    WRITE_STRING(Param->UpdateFileName);
    WRITE_STRING(Param->UpdateLogFileName);
    WRITE_STRING(Param->DownloadPage);

    fclose(File);
}

wchar_t* GetFileNameOffsetW(wchar_t* FullName)
{
    int len = wcslen(FullName);

    while (len--)
    {
        if ((FullName[len] == L'/' ) || (FullName[len] == L'\\' ))
        {
            return FullName + len + 1;
        }
    }

    return FullName;
}

DWORD WINAPI SelfUpdateThread(void* Param)
{
    DWORD Len;
    char* Str;
    TCHAR SelfUpdateProducer[32];
    wchar_t AuExeFileName[MAX_PATH] = {0};
    SELFUPDATEPARAM* P = (SELFUPDATEPARAM*)Param;

    if (GetEnvironmentVariable(TEXT("SelfUpdateProducer"), SelfUpdateProducer, 32))
    {
        ProduceUpdateInfo(P);
    }

    Str = GetContentFromUrl(P->InfoUrl, FALSE, &Len);

    CheckUpdateInfo(P, Str, Len);

    free(Str);

    if (P->UpdateFound)
    {
        if (P->AnswerAllowUpdate == surUndetermined)
        {
            if (P->OnUpdateFound)
                P->OnUpdateFound(P);
            else
                P->AnswerAllowUpdate = surYes;
        }

        switch (P->AnswerAllowUpdate)
        {
        case surYes:
            {
                wchar_t UpdateFullName[MAX_PATH];

                wcscpy(UpdateFullName, P->UpdateFilePath);
                wcscat(UpdateFullName, P->UpdateFileName);

                Str = GetContentFromUrl(UpdateFullName, FALSE, &Len);

                if (Str)
                {
                    HANDLE file;

                    GetModuleFileNameW(NULL, AuExeFileName, MAX_PATH);

                    wcscpy(AuExeFileName + wcslen(AuExeFileName)-4, L"-autoupdated.exe");

                    file = CreateFileW(AuExeFileName, GENERIC_WRITE, FILE_SHARE_READ, NULL,
                        CREATE_ALWAYS, FILE_ATTRIBUTE_NORMAL, NULL);

                    if (file == INVALID_HANDLE_VALUE)
                    {
                        file = NULL;
                    }	

                    if (file)
                    {
                        DWORD nRW = 0;
                        WriteFile(file, Str, Len, &nRW, NULL);

                        CloseHandle(file);
                    }

                    free(Str);

                    if (P->AnswerExitRunningApp == surUndetermined)
                    {
                        if (P->OnDownloaded)
                            P->OnDownloaded(P);
                        else
                            P->AnswerExitRunningApp = surNo;
                    }

                    if (P->AnswerExitRunningApp == surYes)
                    {
                        if (AuExeFileName[0])
                        {
                            ShellExecuteW(NULL, L"open", AuExeFileName, NULL, NULL, SW_NORMAL);
                        }
                        else
                        {
                            wchar_t ExeFileName[MAX_PATH];

                            GetModuleFileNameW(NULL, ExeFileName, MAX_PATH);

                            GetFileNameOffsetW(ExeFileName)[0] = 0;

                            ShellExecuteW(NULL, L"open", ExeFileName, NULL, NULL, SW_NORMAL);

                        }

                        exit(0);
                    }
                }
                else
                {
                    ShellExecuteW(NULL, L"open", P->Homepage, NULL, NULL, SW_NORMAL);
                }

            }
            break;
        case surNo:
            ShellExecuteW(NULL, L"open", P->DownloadPage, NULL, NULL, SW_NORMAL);
            break;
        default:
            break;
        }             

    }
    else
    {
        if (P->OnUpdateNotFound)
        {
            P->OnUpdateNotFound(P);
        }
    }

    free(P);

    return 0;
}

static void OnDownloadedDef(SELFUPDATEPARAM* Param)
{
    if (IDYES == MessageBoxA(NULL, "Download Completed!"
        "\nExit to perform auto updating now?",
        "Notice", MB_ICONQUESTION|MB_YESNO))
    {
        Param->AnswerExitRunningApp = surYes;
    }
    else
    {
        Param->AnswerExitRunningApp = surNo;

        MessageBoxA(NULL, 
            "The auto updating will be performed on next launch of the application",
            "Notice", 0);
    }
}

VOID CALLBACK OnUpdateFoundHelpProc(LPHELPINFO lpHelpInfo)
{
    SELFUPDATEPARAM* Param = (SELFUPDATEPARAM*)lpHelpInfo->dwContextId;
    ShellExecuteW(NULL, L"open", Param->UpdateLogFileName, NULL, NULL, SW_NORMAL);
}

static void OnUpdateFoundDef(SELFUPDATEPARAM* Param)
{
    int ans;
    MSGBOXPARAMSW mbp;
    ZeroMemory(&mbp, sizeof(MSGBOXPARAMS));

    mbp.dwContextHelpId = (DWORD)Param;
    mbp.cbSize = sizeof(mbp);
    mbp.dwStyle = MB_ICONQUESTION|MB_YESNOCANCEL|MB_DEFBUTTON1|MB_HELP;
    mbp.lpfnMsgBoxCallback = OnUpdateFoundHelpProc;
    mbp.lpszCaption = L"Notice";
    mbp.lpszText = L"New Verision Found!\nUpdate to New Version now?\n"
        L"\nYES\tAuto Update in background."
        L"\nNo\tGoto Download Page for downlaoding manually."
        L"\nCancel\tDo not update."
        L"\nHelp\tView the Update Log.";

#if 1
    ans = MessageBoxIndirectW(&mbp);
#else

    int ans = MessageBoxA(NULL, "New Verision Found!\nUpdate to New Version now?\n"
        "\nYES\tAuto Update in background."
        "\nNo\tGoto Download Page for downlaoding manually."
        "\nCancel\tDo not update."
        "\nHelp\tView the Update Log."
        , "Notice", MB_ICONQUESTION|MB_YESNOCANCEL|MB_DEFBUTTON1|MB_HELP);
#endif

    switch (ans)
    {
    case IDYES:
        Param->AnswerAllowUpdate = surYes;
        break;
    case IDNO:
        Param->AnswerAllowUpdate = surNo;
        break;
    default:
        Param->AnswerAllowUpdate = surCancel;
        break;
    }
}

void SelfUpdate(SELFUPDATEPARAM* Param)
{	
    wchar_t RunningFileName[MAX_PATH];

    GetModuleFileNameW(NULL, RunningFileName, MAX_PATH);

    if (_wcsicmp(GetFileNameOffsetW(RunningFileName), Param->AppName) == 0)
    {
        BOOL Updated = TRUE;
        wchar_t TargetFileName[MAX_PATH];

        wcscpy(TargetFileName, RunningFileName);

        wcscpy(TargetFileName + wcslen(TargetFileName)-4, L"-autoupdated.exe");

        if (_waccess(TargetFileName, 0) == 0)
        {
            DWORD tLen;
            DWORD rLen;
            char* Target = GetContentFromUrl(TargetFileName, FALSE, &tLen);
            char* Running = GetContentFromUrl(RunningFileName, FALSE, &rLen);

            if ((rLen != tLen) || (memcmp(Target, Running, tLen) != 0))
            {
                Updated = FALSE;

                ShellExecuteW(NULL, L"open", TargetFileName, NULL, NULL, SW_NORMAL);

                exit(0);
            }
            else
            {
                free(Target);
                free(Running);

                while (_waccess(TargetFileName, 0) == 0)
                {
                    Sleep(300);
                    DeleteFileW(TargetFileName);
                }

                if (Param->UpdateLogFileName[0])
                {
                    if (IDYES == MessageBoxA(NULL, "Already updated to the latest version!\n"
                        "Do you want to view the update log?", "Hint", MB_YESNO|MB_SYSTEMMODAL))
                    {
                        ShellExecuteW(NULL, L"open", Param->UpdateLogFileName, NULL, NULL, SW_NORMAL);
                    }
                }
                else
                {
                    MessageBoxA(NULL, "Update Done!", "Hint", 0);
                }

            }
        }

        if (Updated)
        {
            SELFUPDATEPARAM* ParamDump = (SELFUPDATEPARAM*)malloc(sizeof(SELFUPDATEPARAM));
            DWORD id;

            memcpy(ParamDump, Param, sizeof(SELFUPDATEPARAM));

            if (ParamDump->OnUpdateFound == NULL)
            {
                ParamDump->OnUpdateFound = OnUpdateFoundDef;
            }

            if (ParamDump->OnDownloaded == NULL)
            {
                ParamDump->OnDownloaded = OnDownloadedDef;
            }

            CreateThread(NULL, 0, SelfUpdateThread, ParamDump, 0, &id);

            wcscpy(GetFileNameOffsetW(TargetFileName), Param->UpdateFileName);

            while (_waccess(TargetFileName, 0) == 0)
            {
                Sleep(300);
                DeleteFileW(TargetFileName);
            }	

        }

    }
    else
    {
        wchar_t TargetFileName[MAX_PATH];

        wcscpy(TargetFileName, RunningFileName);

        wcscpy(GetFileNameOffsetW(TargetFileName), Param->AppName);

        while (_waccess(TargetFileName, 0) == 0)
        {
            Sleep(300);
            DeleteFileW(TargetFileName);
        }		

        CopyFileW(RunningFileName, TargetFileName, FALSE);

        ShellExecuteW(NULL, L"open", TargetFileName, NULL, NULL, SW_NORMAL);

        exit(0);
    }
}
